from numpy import array
from sklearn.cluster import KMeans

#训练数据
X = array([[1,1,1,1,1,1,1],
           [2,3,2,2,2,2,2],
           [3,2,3,3,3,3,3],
           [1,2,1,2,2,1,2],
           [2,1,3,3,3,2,1],
           [6,2,30,3,33,2,71]])
#建模
kmeansPredicter = KMeans(n_clusters=3).fit(X)
#训练数据分类
category = kmeansPredicter.predict(X)
print('训练数据分类情况：', category)
#测试
print('1'+'='*30)
X_test=[[1,2,3,3,1,3,1]]
result = kmeansPredicter.predict(X_test)
print('预测结果：', result)
print('相似元素：\n', X[category==result])
print('2'+'='*30)
X_test=[[5,2,23,2,21,5,51]]
result = kmeansPredicter.predict(X_test)
print('预测结果：', result)
print('相似元素：\n', X[category==result])